#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
HTTP数据查询API路由
"""

import logging
import json
from datetime import datetime
from bson import ObjectId
from flask import Blueprint, request, jsonify
from db.mongo import get_db
from intercept.bpf.http_capture import start_capture_for_session, stop_capture

# 自定义JSON编码器，处理MongoDB的ObjectId
class JSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, ObjectId):
            return str(obj)
        if isinstance(obj, datetime):
            return obj.isoformat()
        return json.JSONEncoder.default(self, obj)

logger = logging.getLogger(__name__)
http_data_bp = Blueprint('http_data', __name__)

@http_data_bp.route('/requests', methods=['GET'])
def get_requests():
    """获取HTTP请求数据"""
    session_id = request.args.get('session_id')
    service_name = request.args.get('service_name')

    if not session_id:
        return jsonify({'error': 'Missing session_id parameter'}), 400

    try:
        db = get_db()
        query = {'session_id': session_id}

        if service_name:
            query['service_name'] = service_name

        requests = list(db.requests.find(query, {'_id': 0}))
        return jsonify(requests), 200
    except Exception as e:
        logger.error(f"获取请求数据失败: {e}")
        return jsonify({'error': str(e)}), 500

@http_data_bp.route('/responses', methods=['GET'])
def get_responses():
    """获取HTTP响应数据"""
    session_id = request.args.get('session_id')
    service_name = request.args.get('service_name')
    request_id = request.args.get('request_id')

    if not session_id:
        return jsonify({'error': 'Missing session_id parameter'}), 400

    try:
        db = get_db()
        query = {'session_id': session_id}

        if service_name:
            query['service_name'] = service_name

        if request_id:
            query['request_id'] = request_id

        responses = list(db.responses.find(query, {'_id': 0}))
        return jsonify(responses), 200
    except Exception as e:
        logger.error(f"获取响应数据失败: {e}")
        return jsonify({'error': str(e)}), 500

@http_data_bp.route('/pairs', methods=['GET'])
def get_request_response_pairs():
    """获取请求-响应对"""
    session_id = request.args.get('session_id')
    service_name = request.args.get('service_name')

    if not session_id:
        return jsonify({'error': 'Missing session_id parameter'}), 400

    try:
        db = get_db()
        query = {'session_id': session_id}

        if service_name:
            query['service_name'] = service_name

        # 获取所有请求
        requests = list(db.requests.find(query, {'_id': 0}))

        # 对每个请求查找对应的响应
        pairs = []
        for req in requests:
            resp = db.responses.find_one({
                'session_id': session_id,
                'request_id': req['request_id']
            }, {'_id': 0})

            pairs.append({
                'request': req,
                'response': resp
            })

        return json.dumps(pairs, cls=JSONEncoder), 200, {'Content-Type': 'application/json'}
    except Exception as e:
        logger.error(f"获取请求-响应对失败: {e}")
        return jsonify({'error': str(e)}), 500

@http_data_bp.route('/capture/<session_id>', methods=['POST'])
def start_capture(session_id):
    """启动HTTP捕获"""
    try:
        # 检查会话是否存在
        db = get_db()
        session = db.sessions.find_one({'session_id': session_id})
        if not session:
            return jsonify({'error': 'Session not found'}), 404

        # 启动HTTP捕获
        result = start_capture_for_session(session_id)

        if result:
            return jsonify({'message': f'HTTP capture started for session {session_id}'}), 200
        else:
            return jsonify({'error': 'Failed to start HTTP capture'}), 500
    except Exception as e:
        logger.error(f"启动HTTP捕获失败: {e}")
        return jsonify({'error': str(e)}), 500

@http_data_bp.route('/capture/stop', methods=['POST'])
def stop_capture_endpoint():
    """停止HTTP捕获"""
    try:
        # 停止HTTP捕获
        stop_capture()

        return jsonify({'message': 'HTTP capture stopped'}), 200
    except Exception as e:
        logger.error(f"停止HTTP捕获失败: {e}")
        return jsonify({'error': str(e)}), 500
